#!/usr/bin/python
# -*- coding: utf-8 -*-  
"""
@Project : hello 
@file : MySQLPooledExecutor.py
@Author : shenj
@time : 2025/5/22 9:27
@func : 常用工具-大批量操作mysql的工具类
"""
from typing import Optional, Union, List, Dict, Any
import logging
import pymysql
from pymysql.cursors import DictCursor
from dbutils.pooled_db import PooledDB  # 使用成熟的连接池库

from com.cn.for_cdc.common.cdc_conmons import mysql_pro_connections
from com.cn.for_cdc.common.log_helper import configure_logging

# 配置日志记录
configure_logging()

class MySQLPooledExecutor:
    """基于连接池的高效 MySQL 执行器，支持事务管理和批量操作"""

    def __init__(
            self,
            host: str,
            user: str,
            password: str,
            port: int = 3306,
            database: Optional[str] = None,
            pool_size: int = 5,
            max_overflow: int = 10
    ):
        """
        初始化连接池配置
        :param pool_size: 连接池常驻连接数
        :param max_overflow: 最大溢出连接数
        """
        self.pool = PooledDB(
            creator=pymysql,
            mincached=pool_size,
            maxcached=max_overflow,
            host=host,
            user=user,
            password=password,
            port=port,
            database=database,
            cursorclass=DictCursor,
            autocommit=False,  # 手动控制事务提交
            blocking=True,  # 阻塞模式等待连接
            ping=1,  # 每次使用前自动检查连接活性
        )

    def _get_connection(self):
        """从连接池获取连接（建议通过上下文管理器使用）"""
        return self.pool.connection()

    def query(
            self,
            sql: str,
            params: Optional[Union[tuple, list, dict]] = None
    ) -> List[Dict[str, Any]]:
        """
        执行查询语句（支持参数化）
        :return: 结果列表（字典格式）
        """
        with self._get_connection() as conn:
            try:
                with conn.cursor() as cursor:
                    cursor.execute(sql, params)
                    return cursor.fetchall()
            except pymysql.Error as e:
                logging.error(f"Query failed: {e}\nSQL: {sql}\nParams: {params}")
                raise RuntimeError(f"Query error: {e}") from e

    def execute(
            self,
            sql: str,
            params: Optional[Union[tuple, list, dict]] = None,
            commit: bool = True
    ) -> int:
        """
        执行写操作（支持手动提交事务）
        :param commit: 是否立即提交事务
        :return: 受影响行数
        """
        with self._get_connection() as conn:
            try:
                with conn.cursor() as cursor:
                    rowcount = cursor.execute(sql, params)
                    if commit:
                        conn.commit()
                    return rowcount
            except pymysql.Error as e:
                conn.rollback()
                logging.error(f"Execute failed: {e}\nSQL: {sql}\nParams: {params}")
                raise RuntimeError(f"Execute error: {e}") from e

    def executemany(
            self,
            sql: str,
            params_list: List[Union[tuple, list, dict]],
            batch_size: int = 1000,
            commit: bool = True
    ) -> int:
        """
        批量执行写操作（分批次提交）
        :param batch_size: 单批次操作数量，避免内存溢出[4](@ref)
        :return: 总受影响行数
        """
        total_rows = 0
        with self._get_connection() as conn:
            try:
                with conn.cursor() as cursor:
                    for i in range(0, len(params_list), batch_size):
                        batch_params = params_list[i:i + batch_size]
                        total_rows += cursor.executemany(sql, batch_params)
                        if commit:
                            conn.commit()
                return total_rows
            except pymysql.Error as e:
                conn.rollback()
                logging.error(f"Bulk execute failed: {e}\nSQL: {sql}")
                raise RuntimeError(f"Bulk execute error: {e}") from e

    def verify_connection(self) -> bool:
        """验证连接池是否可用（执行简单查询）"""
        try:
            self.query("SELECT 1")
            logging.info("Connection pool is healthy")
            logging.info(self.pool)
            return True
        except Exception as e:
            logging.error(f"Connection pool verification failed: {e}")
            return False


# 使用示例
if __name__ == '__main__':
    bu = 'drlv'
    db_configs = mysql_pro_connections[f'{bu}_pro_job']
    db = MySQLPooledExecutor(**db_configs, pool_size=5, max_overflow=10)

    # 验证连接池
    db.verify_connection()

    # # 事务示例
    # with db._get_connection() as conn:
    #     try:
    #         conn.begin()  # 显式开启事务
    #         db.execute("UPDATE accounts SET balance = balance - 100 WHERE user_id = %s", (1,), commit=False)
    #         db.execute("UPDATE accounts SET balance = balance + 100 WHERE user_id = %s", (2,), commit=False)
    #         conn.commit()  # 手动提交
    #     except Exception:
    #         conn.rollback()
